import os
import imageio
import math
from PIL import Image

class VideoThumbnailer:
    """Class for generating thumbnails from video."""

    def __init__(self, video_path):
        """
        Initialize VideoThumbnailer object.

        Parameters:
            video_path (str): The path to the input video file.
        """
        self.video_path = video_path
        try:
            self.video = imageio.get_reader(video_path)
            self.num_frames = math.floor(self.video.count_frames())
            self.frame_rate = math.floor(self.video.get_meta_data()['fps'])
        except Exception as e:
            print(f"Error loading video: {str(e)}")
            self.video = None

    def generate_thumbnail(self, frame_num, save_path=None):
        """
        Generate a thumbnail from the video.

        Parameters:
            frame_num (int): The frame number from which to generate the thumbnail.
            save_path (str, optional): The path to save the generated thumbnail. If not provided, a default filename will be used.

        Returns:
            str: The path to the generated thumbnail.
        """
        # Make sure the video was loaded successfully
        if not self.video:
            print("Error generating thumbnail: video not loaded")
            return None

        # Make sure the requested frame is within the video's frame count
        if frame_num >= self.num_frames:
            print("Error generating thumbnail: frame number out of range")
            return None

        # Set the desired frame count
        frame_pos = math.floor(frame_num * self.frame_rate)

        # Capture the frame
        try:
            frame = self.video.get_data(frame_pos)
        except Exception as e:
            print(f"Error generating thumbnail: {str(e)}")
            return None

        # Set the thumbnail size
        thumbnail_size = (640, 420)

        # Resize the frame to the thumbnail size using Pillow
        frame = Image.fromarray(frame).resize(thumbnail_size)

        # Convert the thumbnail to RGB format
        thumbnail = frame.convert('RGB')

        # Get the output filename
        if save_path:
            output_filename = save_path
        else:
            base_filename = os.path.splitext(os.path.basename(self.video_path))[0]
            output_filename = f"{base_filename}.jpg"

        # Save the thumbnail to a file
        try:
            thumbnail.save(output_filename, 'JPEG')
        except Exception as e:
            print(f"Error saving thumbnail: {str(e)}")
            return None

        # Return the output filename
        return output_filename
